"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-3
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""

import sys, getopt
from numpy.core.fromnumeric import repeat
import pylynchipsdk as sdk
import numpy as np
import time

sys.path.append("../../../")
import copy
import common.python.common as common
import common.python.bufferpool as bufferpool
from common.python.blockqueue import block_queue
from common.python.callback_data_struct import recv_cb_data
import threading
from collections import deque

g_exit = False
g_mutex = threading.Lock()
total_send = 0
total_recv = 0
g_running_decoders = 0
g_time_start = time.time()
g_time_end = time.time()
m_stop = False


class decoder_attr:
    def __init__(self) -> None:
        self.url = "../../../data/stop_decode_demo/MOT16-09.mp4"
        self.output_path = ""
        self.outnum_min = 1
        self.outnum_max = 10
        self.chan_id = 0
        self.device_id = 0
        self.output_fmt = sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12
        self.scale = sdk.lyn_scale_t.SCALE_NONE
        self.repeat_num = 1


class video_decoder:
    def __init__(self) -> None:
        self.__attr = ""
        self.__ctx = ""
        self.__send_thread = ""
        self.__recv_thread = ""
        self.__demux_hdl = ""
        self.__codec_para = ""
        self.__vdec_hdl = ""
        self.__vdec_attr = sdk.lyn_vdec_attr_t()
        self.__vdec_out_info = ""
        self.__send_queue = block_queue()
        self.__send_num = 0
        self.__recv_num = 0

    def __send_thread_func(self):
        # 设置上下文环境 创建发送stream
        sdk.lyn_set_current_context(self.__ctx)
        send_stream, ret = sdk.lyn_create_stream()
        common.error_check(ret, "lyn_create_stream")
        eos = False
        repeat = self.__attr.repeat_num
        while not eos:
            # 从解封装器读取一个包
            pkt, ret = sdk.lyn_demux_read_packet(self.__demux_hdl)
            if ret != 0:
                sdk.lyn_demux_close(self.__demux_hdl)
                repeat = repeat - 1
                if repeat > 0:
                    self.__demux_hdl, ret = sdk.lyn_demux_open(self.__attr.url)
                    continue
                eos = True

            # if(self.__send_num == 200): eos = True
            if m_stop == True:
                eos = True
            pkt.eos = eos

            # 发送给解码器解码
            ret = sdk.lyn_vdec_send_packet_async(send_stream, self.__vdec_hdl, pkt)
            common.error_check(ret, "lyn_vdec_send_packet_async")
            ret = sdk.lyn_synchronize_stream(send_stream)
            common.error_check(ret, "lyn_synchronize_stream")

            # 释放packet内存并通知接收结果
            if not eos:
                sdk.lyn_demux_free_packet(pkt)
                self.__send_num = self.__send_num + 1
                self.__send_queue.put(self.__send_num)
                global g_mutex
                with g_mutex:
                    global total_send
                    total_send = total_send + 1
            else:
                self.__send_queue.put(-1)
                if m_stop == True: break
        ret = sdk.lyn_synchronize_stream(send_stream)
        common.error_check(ret, "lyn_synchronize_stream")
        ret = sdk.lyn_destroy_stream(send_stream)
        common.error_check(ret, "lyn_destroy_stream")

    def __recv_thread_func(self):
        # 设置上下文环境 创建接收stream
        sdk.lyn_set_current_context(self.__ctx)
        recv_stream, ret = sdk.lyn_create_stream()

        ipe_stream, ret = sdk.lyn_create_stream()
        vdec_event, ret = sdk.lyn_create_event()

        common.error_check(ret, "lyn_create_stream")
        frame_pool = bufferpool.buffer_pool(self.__vdec_out_info.predict_buf_size, 5)
        while self.__recv_num >= 0:
            self.__recv_num = self.__send_queue.take()
            cb_data = recv_cb_data()
            cb_data.frame.eos = False
            cb_data.frame.data = frame_pool.pop()
            cb_data.frame.size = self.__vdec_out_info.predict_buf_size
            cb_data.frame_pool = frame_pool
            cb_data.send_num = self.__send_num
            cb_data.recv_num = self.__recv_num
            cb_data.attr = self.__attr

            # 插入接收指令，并添加接收完成回调函数
            ret = sdk.lyn_vdec_recv_frame_async(
                recv_stream, self.__vdec_hdl, cb_data.frame
            )
            common.error_check(ret, "lyn_vdec_recv_frame_async")

            ret = sdk.lyn_record_event(recv_stream, vdec_event)
            ret = sdk.lyn_stream_wait_event(ipe_stream, vdec_event)

            ret = sdk.lyn_stream_add_callback(
                ipe_stream, recv_frame_cb, [cb_data]
            )
            common.error_check(ret, "lyn_stream_add_callback")

        ret = sdk.lyn_synchronize_stream(recv_stream)
        common.error_check(ret, "lyn_synchronize_stream")

        ret = sdk.lyn_synchronize_stream(ipe_stream)

        ret = sdk.lyn_destroy_stream(recv_stream)
        common.error_check(ret, "lyn_destroy_stream")

    def run(self, attr: decoder_attr):
        # 创建上下文环境
        self.__attr = copy.copy(attr)
        self.__ctx, ret = sdk.lyn_create_context(attr.device_id)
        common.error_check(ret, "create context")
        ret = sdk.lyn_register_error_handler(common.default_stream_error_handle)
        common.error_check(ret, "lyn_register_error_handler")

        # 打开解封装器
        self.__demux_hdl, ret = sdk.lyn_demux_open(attr.url)
        common.error_check(ret, "lyn_demux_open ")
        self.__codec_para, ret = sdk.lyn_demux_get_codec_para(self.__demux_hdl)
        common.error_check(ret, "lyn_demux_get_codec_para ")

        # 打开解码器
        self.__vdec_attr.codec_id = self.__codec_para.codec_id
        self.__vdec_attr.output_fmt = self.__attr.output_fmt
        self.__vdec_attr.scale = self.__attr.scale
        self.__vdec_hdl, ret = sdk.lyn_vdec_open(self.__vdec_attr)
        common.error_check(ret, "lyn_vdec_open ")
        self.__vdec_out_info, ret = sdk.lyn_vdec_get_out_info(
            self.__codec_para, self.__vdec_attr
        )
        common.error_check(ret, "lyn_vdec_get_out_info ")

        # 开启发送线程
        self.__send_thread = threading.Thread(target=self.__send_thread_func)
        self.__send_thread.start()

        # 开启接收线程
        self.__recv_thread = threading.Thread(target=self.__recv_thread_func)
        self.__recv_thread.start()

    def close(self):
        if self.__send_thread.is_alive():
            self.__send_thread.join()
        if self.__recv_thread.is_alive():
            self.__recv_thread.join()
        if self.__vdec_hdl != "":
            ret = sdk.lyn_vdec_close(self.__vdec_hdl)
            common.error_check(ret, "lyn_vdec_close")
        if self.__ctx != "":
            ret = sdk.lyn_destroy_context(self.__ctx)
            common.error_check(ret, "lyn_destroy_context")


def recv_frame_cb(cb_data: recv_cb_data):
    # 将内存归还给内存池
    para = cb_data[0]
    para.frame_pool.push(para.frame.data)
    global g_mutex
    with g_mutex:
        # 从收到第一帧开始计时，到所有解码器完成结束计时
        if para.recv_num > 0:
            global total_recv
            total_recv = total_recv + 1
            if total_recv == 1:
                global g_time_start
                g_time_start = time.time()
        else:
            global g_running_decoders
            g_running_decoders = g_running_decoders - 1
            if g_running_decoders == 0:
                global g_time_end
                g_time_end = time.time()

        # 保存文件
        if (
            para.attr.output_path != ""
            and para.recv_num >= para.attr.outnum_min
            and para.recv_num <= para.attr.outnum_max
        ):
            host_buf_arr = np.ones(para.frame.size, dtype=np.byte)
            host_buf = sdk.lyn_numpy_to_ptr(host_buf_arr)
            ret = sdk.lyn_memcpy(
                host_buf,
                para.frame.data,
                para.frame.size,
                sdk.lyn_memcpy_dir_t.ServerToClient,
            )
            common.error_check(ret, "lyn_memcpy")
            output_file = "{}dump_{:d}_{:0>4d}.yuv".format(
                para.attr.output_path, para.attr.chan_id, para.recv_num
            )
            with open(output_file, "wb") as f:
                f.write(host_buf_arr.tobytes())
                f.close()
    return 0


def running_show():
    global g_exit
    while not g_exit:
        time.sleep(1 / 100)
        # print("send:", total_send, " / recv", total_recv, "\r", end="", flush=True)


def usage():
    print("usage: python3 main.py")
    print(
        "\t -i    input file, default: /usr/local/lynxi/sdk/sdk-samples/data/MOT16-09.mp4"
    )
    print("\t -o    output path, eg: ./output/")
    print("\t -n    output frame num, default 1,10 eg: num_min,num_max")
    print("\t -d    device id, default 0")
    print("\t -c    channel count, default 1")
    print("\t -r    repeat count, default 1")
    print("\t -h    show help")
    print("\t -?    show help")


def main(argv):
    try:
        opts, args = getopt.getopt(argv, "i:o:n:d:c:r:h?")
    except getopt.GetoptError:
        usage()
        return 0

    attr = decoder_attr()
    channel_count = 1

    for opt, arg in opts:
        if opt in ("-h", "-?"):
            usage()
            return 0
        elif opt == "-i":
            attr.url = arg
        elif opt == "-c":
            channel_count = int(arg)
        elif opt == "-r":
            attr.repeat_num = int(arg)
        elif opt == "-o":
            attr.output_path = arg
        elif opt == "-n":
            num = arg.split(",")
            attr.outnum_min = int(num[0])
            attr.outnum_max = int(num[1])
        elif opt == "-d":
            attr.device_id = int(arg)
        else:
            usage()
            return 0

    if attr.url == "":
        print("input file miss!!!")
        usage()
        return 0

    running_thread = threading.Thread(target=running_show)
    running_thread.start()

    global g_running_decoders
    g_running_decoders = channel_count

    decoders = deque()
    for i in range(channel_count):
        attr.chan_id = i
        decoder = video_decoder()
        decoder.run(attr)
        decoders.append(decoder)

    name = input("Press 'Enter' key to exit: ")
    global m_stop
    m_stop = True

    for it in decoders:
        it.close()

    global g_exit
    g_exit = True
    running_thread.join()
    global g_time_start
    global g_time_end
    frame_rate = "{:.2f}".format(total_recv / (g_time_end - g_time_start))
    print()
    print("decode num ", total_recv, "frame_rate ", frame_rate, " fps")


if __name__ == "__main__":
    main(sys.argv[1:])
